import asyncio
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

from pytdx.hq import TdxHq_API

# 创建API对象
api = TdxHq_API()

# 连接到通达信服务器（示例IP和端口，实际使用时需替换为有效的服务器地址和端口）
# 124.70.133.119  47.116.105.28 121.36.54.217 121.36.81.195 123.249.15.60
def get_data(stock_code, market_type = 0, line_type=9, start=0, count=800):
    with api.connect('121.36.81.195', 7709):
        # 获取股票日线数据
        # 参数解释：市场代码（0：深圳，1：上海），股票代码，K线种类（9：日线），开始位置，请求数量
        data = api.to_df(api.get_security_bars(line_type, market_type, stock_code, start, count))
        # 打印获取的数据
        return data
    
def strategy_mv(data):
    # 将日期列转换为日期格式
    data['Date'] = pd.to_datetime(data['datetime'])

    # 设置日期列为索引
    data.set_index('Date', inplace=True)

    # 计算短期和长期移动平均线
    short_window = 20
    long_window = 50
    data['short_mavg'] = data['close'].rolling(window=short_window).mean()
    data['long_mavg'] = data['close'].rolling(window=long_window).mean()

    # 生成交易信号
    data['signal'] = np.where(data['short_mavg'] > data['long_mavg'], 1, 0)
    data['positions'] = data['signal'].diff()

    # 计算策略的每日收益率
    data['returns'] = data['close'].pct_change()
    data['strategy_returns'] = data['returns'] * data['positions'].shift(1)

    # 计算累计收益率
    data['cumulative_returns'] = (1 + data['strategy_returns']).cumprod()

    return data

def show_chart(data):
    plt.close('all')
    # 绘制股票价格和移动平均线
    plt.figure(figsize=(10, 6))
    plt.plot(data['close'], label='close Price')
    plt.plot(data['short_mavg'], label='Short Moving Average')
    plt.plot(data['long_mavg'], label='Long Moving Average')
    plt.title('Stock Price and Moving Averages')
    plt.xlabel('Date')
    plt.ylabel('Price')
    plt.legend()
    plt.show()

    # 绘制策略的累计收益率
    plt.figure(figsize=(10, 6))
    plt.plot(data['cumulative_returns'], label='Strategy Cumulative Returns')
    plt.title('Strategy Cumulative Returns')
    plt.xlabel('Date')
    plt.ylabel('Cumulative Returns')
    plt.legend()
    plt.show()

# 定义事件类
class MarketDataEvent:
    def __init__(self, data):
        self.data = data
        print(type(self.data))
    
    def push_data(self, new_data):
        self.data = pd.concat([self.data, new_data])

# 定义事件处理函数
class Strategy:
    def __init__(self):
        self.short_window = 20
        self.long_window = 50
        self.signals = []

    async def on_market_data(self, event):
        data = event.data
        event.data = strategy_mv(data)


async def event_loop(strategy):
    data = get_data('000001')
    event = MarketDataEvent(data[:1])
    for i in range(1, len(data)):
        event.push_data(data[i:i+1])
        await strategy.on_market_data(event)
    show_chart(event.data)

async def main():
    strategy = Strategy()
    await event_loop(strategy)

asyncio.run(main())


